import os
import random
import xml.etree.ElementTree as ET


# work_dit = "./data/98"
# print(os.listdir(work_dit))

def saved(var, filepath):
    with open(filepath, 'w') as f:
        var_str = str(var).replace("'", "\"")
        f.write(var_str)


def get_label_map():
    label_map = dict({
        'rat': 1,
        'background': 0
    })
    saved(label_map, './label_map.json')


def get_xml(path):
    boxes = []
    labels = []
    difficulties = []
    tree = ET.parse(path)
    root = tree.getroot()
    for object in root.findall('object'):
        if object[0].text == 'rat':
            labels.append(1)
            difficulties.append(0)
        else:
            labels.append(0)
            difficulties.append(0)
        for box in object.findall('bndbox'):
            boxes.append([int(box[0].text), int(box[1].text), int(box[2].text), int(box[3].text)])
    return boxes, labels, difficulties


def data_split(full_list, ratio, shuffle=False):
    """
    数据集拆分: 将列表full_list按比例ratio（随机）划分为2个子列表sublist_1与sublist_2
    :param full_list: 数据列表
    :param ratio:     比例
    :param shuffle:   是否存在子列表2
    :return:
    """
    n_total = len(full_list)
    offset = int(n_total * ratio)
    if n_total == 0 or offset < 1:
        return [], full_list
    if shuffle:
        random.shuffle(full_list)
    sublist_1 = full_list[:offset]
    sublist_2 = full_list[offset:]
    return sublist_1, sublist_2


def get_id(work_path):
    work_dir = work_path
    dirlist = os.listdir(work_dir)
    id = []
    for i in dirlist:
        i = os.path.splitext(i)
        id.append(i[0])
    id = list(set(id))
    return id


def get_images(work_path):
    get_label_map()
    id = get_id(work_path)
    test_id, train_id = data_split(id, ratio=0.3, shuffle=True)
    test_images, train_images = [], []
    test_object, train_object = [], []
    for i in test_id:
        di = {}
        boxes, labels, difficulties = get_xml(work_path + '/' + i + '.xml')
        if not boxes:
            continue
        test_images.append(os.path.abspath(work_path + '/' + i + '.jpg'))
        di['boxes'] = boxes
        di['labels'] = labels
        di['difficulties'] = difficulties
        test_object.append(di)
    for j in train_id:
        dj = {}
        boxes, labels, difficulties = get_xml(work_path + '/' + j + '.xml')
        if not boxes:
            continue
        train_images.append(os.path.abspath(work_path + '/' + j + '.jpg'))
        dj['boxes'] = boxes
        dj['labels'] = labels
        dj['difficulties'] = difficulties
        train_object.append(dj)
    saved(test_object, './TEST_objects.json')
    saved(train_object, './TRAIN_objects.json')
    saved(test_images, './TEST_images.json')
    saved(train_images, './TRAIN_images.json')


if __name__ == '__main__':
    get_images('./data/98')